package leetcode.lcp0046;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public class Solution {
	public int[] volunteerDeployment(int[] finalCnt, long totalNum, int[][] edges, int[][] plans) {
		int n = finalCnt.length;
		List<Set<Integer>> nexts = new ArrayList<>();
		for (int i = 0; i <= n; i++) {
			nexts.add(new HashSet<>());
		}
		for (int[] e : edges) {
			int x = e[0], y = e[1];
			nexts.get(x).add(y);
			nexts.get(y).add(x);
		}
		int[] ans = new int[n + 1];
		ans[0] = 1;
		long[] cnt = new long[n + 1];
		for (int i = 0; i < n; i++) {
			cnt[i + 1] = finalCnt[i];
		}
		for (int i = plans.length - 1; i >= 0; i--) {
			int[] plan = plans[i];
			int num = plan[0], idx = plan[1];
			if (num == 1) {
				// idx场馆人数翻倍
				f1(idx, ans, cnt);
			} else if (num == 2) {
				// idx相邻减去idx数量
				f2(idx, nexts, ans, cnt);
			} else { // num ==3
				// idx相邻加上idx数量
				f3(idx, nexts, ans, cnt);
			}
		}
		for (long num : cnt) {
			totalNum -= num;
		}
		int sumAns = Arrays.stream(ans).sum();
		long avg = totalNum / sumAns;
		for (int i = 0; i <= n; i++) {
			ans[i] = (int) (ans[i] * avg + cnt[i]);
		}
		return ans;
	}

	private void f1(int idx, int[] ans, long[] cnt) {
		ans[idx] = ans[idx] << 1;
		cnt[idx] = cnt[idx] << 1;
	}

	private void f3(int idx, List<Set<Integer>> nexts, int[] ans, long[] cnt) {
		Set<Integer> next = nexts.get(idx);
		for (int n : next) {
			ans[n] += ans[idx];
			cnt[n] += cnt[idx];
		}
	}

	private void f2(int idx, List<Set<Integer>> nexts, int[] ans, long[] cnt) {
		Set<Integer> next = nexts.get(idx);
		for (int n : next) {
			ans[n] -= ans[idx];
			cnt[n] -= cnt[idx];
		}
	}
}
